#!/usr/bin/env python3
"""
PDF转TXT工具
使用PyMuPDF (fitz)库将PDF文件转换为纯文本格式

作者: AutoNovel项目
日期: 2025-10-17
"""

import os
import sys
import argparse
from pathlib import Path
from typing import Optional

try:
    import fitz  # PyMuPDF
except ImportError:
    print("错误: 请先安装PyMuPDF库")
    print("安装命令: uv add pymupdf 或 pip install pymupdf")
    sys.exit(1)


class PDFConverter:
    """PDF转换器类"""

    def __init__(self, encoding: str = 'utf-8'):
        """
        初始化转换器

        Args:
            encoding: 输出文本编码格式，默认为utf-8
        """
        self.encoding = encoding

    def detect_and_fix_encoding(self, text: str) -> str:
        """
        检测和修复文本编码问题

        Args:
            text: 可能存在编码问题的文本

        Returns:
            修复后的文本
        """
        if not text:
            return text

        # 常见的中文编码
        encodings_to_try = ['utf-8', 'gbk', 'gb2312', 'big5', 'latin1', 'cp1252']

        for encoding in encodings_to_try:
            try:
                # 尝试编码再解码来清理文本
                if encoding == 'utf-8':
                    # UTF-8文本，直接返回
                    return text
                else:
                    # 尝试其他编码
                    cleaned_text = text.encode(encoding, errors='ignore').decode('utf-8', errors='ignore')
                    if cleaned_text and len(cleaned_text.strip()) > 10:
                        return cleaned_text
            except (UnicodeEncodeError, UnicodeDecodeError):
                continue

        # 如果所有编码都失败，尝试使用chardet（如果可用）
        try:
            import chardet
            detected = chardet.detect(text.encode('latin1', errors='ignore'))
            if detected and detected['confidence'] > 0.7:
                encoding = detected['encoding']
                return text.encode('latin1', errors='ignore').decode(encoding, errors='ignore')
        except ImportError:
            pass
        except:
            pass

        # 最后的备选方案：清理非打印字符
        import re
        # 保留中文字符、英文字符、数字和常用标点
        cleaned_text = re.sub(r'[^\x20-\x7E\u4e00-\u9fff\u3000-\u303f\uff00-\uffef]', '', text)
        return cleaned_text

    def extract_text_from_page(self, page, clean: bool = True) -> str:
        """
        从单个页面提取文本

        Args:
            page: PyMuPDF页面对象
            clean: 是否清理空白字符

        Returns:
            提取的文本内容
        """
        # 尝试多种文本提取方法
        text = None

        # 方法1: 使用get_text()的默认参数
        text = page.get_text()

        # 如果结果为空或包含大量乱码，尝试其他方法
        if not text or len(text.strip()) < 10:
            # 方法2: 使用不同的提取参数
            text = page.get_text("text")

        if not text or len(text.strip()) < 10:
            # 方法3: 使用HTML格式然后提取文本
            html_text = page.get_text("html")
            if html_text and len(html_text.strip()) > 10:
                # 简单去除HTML标签
                import re
                text = re.sub(r'<[^>]+>', '', html_text)

        if not text or len(text.strip()) < 10:
            # 方法4: 使用字典格式提取
            try:
                blocks = page.get_text("dict")
                text_parts = []
                for block in blocks["blocks"]:
                    if "lines" in block:
                        for line in block["lines"]:
                            for span in line["spans"]:
                                text_parts.append(span["text"])
                text = "".join(text_parts)
            except:
                pass

        # 最后的备选方案：使用图像识别（如果安装了相关依赖）
        if not text or len(text.strip()) < 10:
            try:
                text = page.get_text("words")
                if text:
                    words = [w[4] for w in text]  # 提取文本部分
                    text = " ".join(words)
            except:
                pass

        # 应用编码修复
        if text:
            text = self.detect_and_fix_encoding(text)

        if clean and text:
            # 清理多余的空白字符
            lines = text.split('\n')
            cleaned_lines = []
            for line in lines:
                stripped_line = line.strip()
                if stripped_line:  # 只保留非空行
                    cleaned_lines.append(stripped_line)
            text = '\n'.join(cleaned_lines)

        return text or ""

    def convert_pdf_to_txt(self, pdf_path: str, output_path: Optional[str] = None,
                          clean: bool = True, debug: bool = False) -> str:
        """
        将PDF文件转换为TXT文件

        Args:
            pdf_path: PDF文件路径
            output_path: 输出TXT文件路径，如果为None则自动生成
            clean: 是否清理空白字符
            debug: 是否启用调试模式，显示详细转换信息

        Returns:
            输出文件路径
        """
        pdf_path = Path(pdf_path)

        if not pdf_path.exists():
            raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")

        if pdf_path.suffix.lower() != '.pdf':
            raise ValueError("输入文件必须是PDF格式")

        # 生成输出文件路径
        if output_path is None:
            output_path = pdf_path.with_suffix('.txt')
        else:
            output_path = Path(output_path)

        print(f"正在转换: {pdf_path}")
        print(f"输出文件: {output_path}")

        try:
            # 打开PDF文件
            doc = fitz.open(str(pdf_path))

            # 提取所有页面的文本
            full_text = []
            total_pages = doc.page_count

            print(f"PDF共有 {total_pages} 页")

            for page_num in range(total_pages):
                page = doc[page_num]
                page_text = self.extract_text_from_page(page, clean)

                if debug:
                    print(f"\n=== 第 {page_num + 1} 页调试信息 ===")
                    print(f"提取的文本长度: {len(page_text)}")
                    print(f"文本预览: {page_text[:100]}...")

                # 添加页码标记（可选）
                if total_pages > 1:
                    full_text.append(f"\n{'='*50}")
                    full_text.append(f"第 {page_num + 1} 页")
                    full_text.append('='*50)

                full_text.append(page_text)

                # 显示进度
                print(f"已处理第 {page_num + 1}/{total_pages} 页", end='\r')

            # 关闭文档
            doc.close()

            # 保存到TXT文件
            final_text = '\n'.join(full_text)

            with open(output_path, 'w', encoding=self.encoding) as f:
                f.write(final_text)

            print(f"\n转换完成! 文件已保存到: {output_path}")
            print(f"文件大小: {output_path.stat().st_size} 字节")

            return str(output_path)

        except Exception as e:
            print(f"转换过程中发生错误: {str(e)}")
            raise

    def batch_convert(self, input_dir: str, output_dir: Optional[str] = None,
                     pattern: str = "*.pdf", clean: bool = True) -> list:
        """
        批量转换PDF文件

        Args:
            input_dir: 输入目录路径
            output_dir: 输出目录路径，如果为None则使用输入目录
            pattern: 文件匹配模式，默认为*.pdf
            clean: 是否清理空白字符

        Returns:
            成功转换的文件路径列表
        """
        input_path = Path(input_dir)

        if not input_path.exists():
            raise FileNotFoundError(f"输入目录不存在: {input_path}")

        if output_dir is None:
            output_path = input_path
        else:
            output_path = Path(output_dir)
            output_path.mkdir(exist_ok=True)

        # 查找所有PDF文件
        pdf_files = list(input_path.glob(pattern))

        if not pdf_files:
            print(f"在目录 {input_path} 中未找到匹配 {pattern} 的PDF文件")
            return []

        print(f"找到 {len(pdf_files)} 个PDF文件")

        converted_files = []

        for pdf_file in pdf_files:
            try:
                output_file = output_path / pdf_file.with_suffix('.txt').name
                result_path = self.convert_pdf_to_txt(pdf_file, output_file, clean)
                converted_files.append(result_path)
                print("-" * 50)
            except Exception as e:
                print(f"转换文件 {pdf_file} 时出错: {str(e)}")
                continue

        print(f"\n批量转换完成! 成功转换 {len(converted_files)}/{len(pdf_files)} 个文件")
        return converted_files


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="PDF转TXT工具 - 将PDF文件转换为纯文本格式",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 转换单个PDF文件
  python pdf_converter.py docs/ch01.pdf

  # 指定输出路径
  python pdf_converter.py docs/ch01.txt docs/ch01_converted.txt

  # 批量转换目录中的所有PDF文件
  python pdf_converter.py --batch docs/

  # 批量转换并指定输出目录
  python pdf_converter.py --batch docs/ output_txt/ --clean
        """
    )

    parser.add_argument('input', nargs='?', help='输入PDF文件路径或目录路径（用于批量转换）')
    parser.add_argument('output', nargs='?', help='输出TXT文件路径（可选）')
    parser.add_argument('--batch', '-b', action='store_true',
                       help='批量转换模式，转换指定目录中的所有PDF文件')
    parser.add_argument('--output-dir', '-o', help='批量转换时的输出目录（可选）')
    parser.add_argument('--pattern', '-p', default='*.pdf',
                       help='批量转换时的文件匹配模式（默认: *.pdf）')
    parser.add_argument('--clean', '-c', action='store_true', default=True,
                       help='清理空白字符（默认启用）')
    parser.add_argument('--no-clean', action='store_true',
                       help='不清理空白字符')
    parser.add_argument('--encoding', '-e', default='utf-8',
                       help='输出文件编码格式（默认: utf-8）')
    parser.add_argument('--debug', '-d', action='store_true',
                       help='启用调试模式，显示详细转换信息')

    args = parser.parse_args()

    # 处理清理选项
    clean = args.clean and not args.no_clean

    # 创建转换器
    converter = PDFConverter(encoding=args.encoding)

    try:
        if args.batch:
            # 批量转换模式
            if not args.input:
                print("错误: 批量转换模式需要指定输入目录")
                sys.exit(1)

            converted_files = converter.batch_convert(
                args.input,
                args.output_dir,
                args.pattern,
                clean
            )

            if converted_files:
                print("\n转换成功的文件:")
                for file_path in converted_files:
                    print(f"  - {file_path}")
            else:
                print("没有文件被成功转换")

        else:
            # 单文件转换模式
            if not args.input:
                print("错误: 请指定要转换的PDF文件")
                parser.print_help()
                sys.exit(1)

            output_path = converter.convert_pdf_to_txt(
                args.input,
                args.output,
                clean,
                args.debug
            )

            print(f"\n转换完成! 输出文件: {output_path}")

    except Exception as e:
        print(f"错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()